import torch
from torch.nn import MSELoss

def mse_map(
    out, gt_map, **kwargs
):
    
    mask = (gt_map > 0.0)
    gt_val = gt_map[mask]
    pred_val = out[mask]

    fn = MSELoss()
    loss = fn(pred_val, gt_val)

    return loss

def peaky_map(
    out, gt_map, **kwargs
):

    loss = 2.0 
    - torch.max(torch.flatten(out, start_dim=1), 1)
    + torch.mean(out, (1, 2))

    return loss
